{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerated Inference With PEFT'd StarCoder2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the previous [notebook](https://github.com/NVIDIA/GenerativeAIExamples/blob/main/models/StarCoder2/lora.ipynb), we show how to parameter efficiently finetune StarCoder2 model with a custom code (instruction, completion) pair dataset. We choose LoRA as our PEFT algorithnm and finetune for 50 interations. In this notebook, the goal is to demonstrate how to compile fintuned .nemo model into optimized TensorRT-LLM engines. The converted model engine can perform accelerated inference locally or be deployed to Triton Inference Server."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export Model Via TensorRT-LLM\n",
    "\n",
    "NVIDIA TensorRT-LLM is an open-source library that accelerates and optimizes inference performance of the latest LLMs on supported AI platforms. NVIDIA NeMo framework offers TensorRT-LLM as an user friendly tool to compile .nemo models into optimized engines. To start with, let's create a folder where the exported model files will be saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir starcoder2_trt_llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to create an instance of the TensorRTLLM class and call the TensorRTLLM.export() function with the nemo_checkpoint_path pointing to the LoRA fine-tuned .nemo checkpoint.\n",
    "\n",
    "After optimized model export, a few files will be stored in the folder we just created. These files include an engine file that holds the weights, the compiled execution graph of the model, a tokenizer.model file which contains the tokenizer information, and config.json which records the metadata about the model (along with model.cache, which caches some operations and makes it faster to re-compile the model in the future.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.export import TensorRTLLM\n",
    "trt_llm_exporter = TensorRTLLM(model_dir=\"starcoder2_trt_llm\")\n",
    "trt_llm_exporter.export(nemo_checkpoint_path=\"starcoder2_lora_alpaca_python_merged.nemo\", model_type=\"starcoder\", n_gpus=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the finetuned model is exported into TensorRT-LLM optimized engines, we can perform accelerated inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trt_llm_exporter.forward([\"Given a non-empty array of integers nums, every element appears twice except for one. Find that single one. ### Input: nums = [2,2,1] ### Output: 1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another code generation example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trt_llm_exporter.forward([\"Implement Fibonacci sequence in Python\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy Model Using Triton Inference Server\n",
    "\n",
    "Lastly, we can easily deploy the finetuned model as a service, which is supported by Triton Inference Server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.deploy import DeployPyTriton\n",
    "\n",
    "nm = DeployPyTriton(model=trt_llm_exporter, triton_model_name=\"starcoder\")\n",
    "nm.deploy()\n",
    "nm.serve()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
